import os
import time
import yaml
import numpy as np
import mayavi.mlab as mlab
import torch
import json
import cv2

# 渲染时不需要显示图片
# mlab.options.offscreen = True

def get_lidar(path_lidar):
    points = np.fromfile(str(path_lidar),dtype=np.float32).reshape(-1,4)
    return points

def check_numpy_to_torch(x):
    if isinstance(x, np.ndarray):
        return torch.from_numpy(x).float(), True
    return x, False

def rotate_points_along_z(points, angle):
    """
    Args:
        points: (B, N, 3 + C)
        angle: (B), angle along z-axis, angle increases x ==> y
    Returns:

    """
    points, is_numpy = check_numpy_to_torch(points)
    angle, _ = check_numpy_to_torch(angle)

    cosa = torch.cos(angle)
    sina = torch.sin(angle)
    zeros = angle.new_zeros(points.shape[0])
    ones = angle.new_ones(points.shape[0])
    rot_matrix = torch.stack((
        cosa, sina, zeros,
        -sina, cosa, zeros,
        zeros, zeros, ones
    ), dim=1).view(-1, 3, 3).float()
    points_rot = torch.matmul(points[:, :, 0:3], rot_matrix)
    points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1)
    return points_rot.numpy() if is_numpy else points_rot


def boxes_to_corners_3d(boxes3d):
    """
        7 -------- 4
       /|         /|
      6 -------- 5 .
      | |        | |
      . 3 -------- 0
      |/         |/
      2 -------- 1
    Args:
        boxes3d:  (N, 7) [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center

    Returns:
    """
    boxes3d, is_numpy = check_numpy_to_torch(boxes3d)

    template = boxes3d.new_tensor((
        [1, 1, -1], [1, -1, -1], [-1, -1, -1], [-1, 1, -1],
        [1, 1, 1], [1, -1, 1], [-1, -1, 1], [-1, 1, 1],
    )) / 2

    corners3d = boxes3d[:, None, 3:6].repeat(1, 8, 1) * template[None, :, :]
    corners3d = rotate_points_along_z(corners3d.view(-1, 8, 3), boxes3d[:, 6]).view(-1, 8, 3)
    corners3d += boxes3d[:, None, 0:3]

    return corners3d.numpy() if is_numpy else corners3d


def visualize_pts(pts, fig=None, bgcolor=(0, 0, 0), fgcolor=(1.0, 1.0, 1.0),color = (0.0, 0.0, 0.5),
                  show_intensity=False, size=(600, 600), draw_origin=True):
    if not isinstance(pts, np.ndarray):
        pts = pts.cpu().numpy()
    if fig is None:
        fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=fgcolor, engine=None, size=size)

    if show_intensity:
        G = mlab.points3d(pts[:, 0], pts[:, 1], pts[:, 2], pts[:, 4]*255, mode='point',
                          colormap='spectral', scale_factor=1, figure=fig)
    else:
        G = mlab.points3d(pts[:, 0], pts[:, 1], pts[:, 2], mode='point',
                          color=color, scale_factor=1, figure=fig)
    if draw_origin:
        mlab.points3d(0, 0, 0, color=(1, 1, 1), mode='cube', scale_factor=0.2)
        mlab.plot3d([0, 3], [0, 0], [0, 0], color=(0, 0, 1), tube_radius=0.1)
        mlab.plot3d([0, 0], [0, 3], [0, 0], color=(0, 1, 0), tube_radius=0.1)
        mlab.plot3d([0, 0], [0, 0], [0, 3], color=(1, 0, 0), tube_radius=0.1)

    for i in range(7):
        mlab.points3d(0, 0, 0, ((i+1)*50), color=(1, 1, 1), mode="2dcircle", scale_factor=2, resolution=100)

    return fig


def draw_sphere_pts(pts, color=(0, 1, 0), fig=None, bgcolor=(0, 0, 0), scale_factor=0.2):
    if not isinstance(pts, np.ndarray):
        pts = pts.cpu().numpy()

    if fig is None:
        fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=None, engine=None, size=(600, 600))

    if isinstance(color, np.ndarray) and color.shape[0] == 1:
        color = color[0]
        color = (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0)

    if isinstance(color, np.ndarray):
        pts_color = np.zeros((pts.__len__(), 4), dtype=np.uint8)
        pts_color[:, 0:3] = color
        pts_color[:, 3] = 255
        G = mlab.points3d(pts[:, 0], pts[:, 1], pts[:, 2], np.arange(0, pts_color.__len__()), mode='sphere',
                          scale_factor=scale_factor, figure=fig)
        G.glyph.color_mode = 'color_by_scalar'
        G.glyph.scale_mode = 'scale_by_vector'
        G.module_manager.scalar_lut_manager.lut.table = pts_color
    else:
        mlab.points3d(pts[:, 0], pts[:, 1], pts[:, 2], mode='sphere', color=color,
                      colormap='gnuplot', scale_factor=scale_factor, figure=fig)

    mlab.points3d(0, 0, 0, color=(1, 1, 1), mode='cube', scale_factor=0.2)
    mlab.plot3d([0, 3], [0, 0], [0, 0], color=(0, 0, 1), line_width=3, tube_radius=None, figure=fig)
    mlab.plot3d([0, 0], [0, 3], [0, 0], color=(0, 1, 0), line_width=3, tube_radius=None, figure=fig)
    mlab.plot3d([0, 0], [0, 0], [0, 3], color=(1, 0, 0), line_width=3, tube_radius=None, figure=fig)

    return fig


def draw_grid(Lx, Ly, Lz, fig, tube_radius=None, color=(0.5, 0.5, 0.5)):
    mlab.plot3d(Lx, Ly, Lz, color=color, tube_radius=tube_radius, line_width=1, figure=fig)
    return fig


def draw_multi_grid_range(fig, grid_size=50, bv_range=(-60, -60, 60, 60)):
    Lx, Ly, Lz, Lx1, Ly1, Lz1 = [], [], [], [], [], []

    Lx = np.arange(bv_range[0], bv_range[2] + 1, grid_size)
    Ly1 = np.arange(bv_range[1], bv_range[3] + 1, grid_size)
    Lx = np.stack([Lx, Lx], axis=0).transpose().reshape(-1)
    Ly1 = np.stack([Ly1, Ly1], axis=0).transpose().reshape(-1)

    for i in range(int(Lx.shape[0] / 2)):
        if (i + 1) % 2 == 0:
            Ly.append(bv_range[1])
            Ly.append(bv_range[1 + 2])
        else:
            Ly.append(bv_range[1 + 2])
            Ly.append(bv_range[1])
    Ly = np.array(Ly)
    Lz = np.zeros(Ly.shape[0])

    for i in range(int(Ly1.shape[0] / 2)):
        if (i + 1) % 2 == 0:
            Lx1.append(bv_range[0])
            Lx1.append(bv_range[0 + 2])
        else:
            Lx1.append(bv_range[0 + 2])
            Lx1.append(bv_range[0])
    # print(Lx)
    Lx1 = np.array(Lx1)
    Lz1 = np.zeros(Lx1.shape[0])

    fig = draw_grid(Lx, Ly, Lz, fig)
    fig = draw_grid(Lx1, Ly1, Lz1, fig)

    # Lx = np.concatenate([Lx,L])
    return fig

def draw_label3d(box=None,name=None,config=None,fig=None,bgcolor=(0, 0, 0), fgcolor=(1.0, 1.0, 1.0),size=[500,600]):
    if fig is None:
        fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=fgcolor, engine=None, size=size)
    # 绘制目标
    colors = [config['color_maps']['classes'][i[0:2]] for i in name]
    if box is not None:
        corners3d = boxes_to_corners_3d(box)
        fig = draw_box3d(corners3d, fig=fig, color=colors ,label=name, max_num=100, box=box)
    types = np.array([i[0:2] for i in name])
    for type in config['color_maps']['classes']:
        # 车辆位置从几何中点变换到车尾中点
        box_temp = box[types == type]
        l = box_temp[:, 3].reshape([1, 1, box_temp.__len__()])

        phi = box_temp[:, 6] * 180 / np.pi
        # phi = phi % 360
        # phi[phi > 270] = phi[phi > 270] - 360
        # phi[(phi > 90) & (phi <= 270)] = phi[(phi > 90) & (phi <= 270)] - 180
        #
        # # 车辆位置从几何中点变换到车位中点
        sin_phi = np.sin(phi / 180 * np.pi)
        cos_phi = np.cos(phi / 180 * np.pi)
        #
        # box_temps = np.array([cos_phi, sin_phi]).T * (-1 / 2) * np.array([box_temp[:, 3], box_temp[:, 3]]).T + box_temp[
        #                                                                                                        :, 0:2]

        x = box_temp[:, 0].reshape([1, 1, box_temp.__len__()])
        y = box_temp[:, 1].reshape([1, 1, box_temp.__len__()])
        z = box_temp[:, 2].reshape([1, 1, box_temp.__len__()])
        u = x / x * cos_phi.reshape([1, 1, box_temp.__len__()]) * l
        v = y / y * sin_phi.reshape([1, 1, box_temp.__len__()]) * l
        w = z
        w[:] = 0
        mlab.quiver3d(x, y, z, u, v, w, figure=fig, color=config['color_maps']['classes'][type], scale_factor=1)
    return fig

def draw_box3d(corners3d, fig, color=(1, 0, 0), line_width=2, label=None, tag='', max_num=500, tube_radius=None, box=None):
    """
    :param corners3d: (N, 8, 3)
    :param fig:
    :param color:
    :param line_width:
    :param cls:
    :param tag:
    :param max_num:
    :return:
    """
    num = min(max_num, len(corners3d))
    for n in range(num):
        b = corners3d[n]  # (8, 3)

        Lx = b[:, 0].T
        Ly = b[:, 1].T
        Lz = b[:, 2].T

        if isinstance(color, tuple):
            colord = color
        else:
            # colord = tuple(color_table_for_class[color[n]])
            colord = color[n]
        if label is not None:
            mlab.text3d(b[6, 0], b[6, 1], b[6, 2], '%s' % label[n], scale=(0.3, 0.3, 0.3), color=colord, figure=fig)

        triangles = [(0, 1, 2), (0, 2, 3), (0, 3, 7), (0, 7, 4), (0, 4, 5), (0, 5, 1), (6, 1, 2), (6, 2, 3), (6, 3, 7),
                     (6, 7, 4), (6, 4, 5), (6, 5, 1)]

        mlab.triangular_mesh(Lx, Ly, Lz, triangles, color=colord, opacity=0.5, figure=fig)
    return fig

def get_json_data(label_file):
    # 读取预测数据
    with open(label_file, 'r') as f:
        data_label = json.load(f)
    name = [i['obj_type'].lower() for i in data_label]
    id = [i['obj_id'] for i in data_label]
    xyz = np.array(
        [[i['psr']['position']['x'], i['psr']['position']['y'], i['psr']['position']['z']] for i in data_label])
    lwh = np.array([[i['psr']['scale']['x'], i['psr']['scale']['y'], i['psr']['scale']['z']] for i in data_label])
    r = np.array([[i['psr']['rotation']['z']] for i in data_label])
    box = np.concatenate([xyz, lwh, r], axis=-1)
    box[:, 2] = box[:, 2] - 3.15
    velocity = np.zeros([name.__len__(), 5]).tolist()
    acceleration = np.zeros([name.__len__(), 5]).tolist()
    if 'gt' not in label_file:
        score = [i['score'] for i in data_label]
        return name, id, box.tolist(), velocity, acceleration, score
    else:
        return name, id, box.tolist(), velocity, acceleration

def get_jsons_data(label_file):
    # 读取检测数据
    with open(label_file, 'r') as f:
        data_label = json.load(f)
    name = data_label['annos']['name']
    id = data_label['annos']['track_id']
    box = data_label['annos']['gt_boxes_3d']
    box = np.array(box)
    velocity = data_label['annos']['velocity']
    acceleration = data_label['annos']['velocity']
    if 'gt' not in label_file:
        score = data_label['annos']['confidence']
        return name, id, box.tolist(), velocity, acceleration, score
    else:
        return name, id, box.tolist(), velocity, acceleration

def resize_keep_aspectratio(image_src, dst_size):
    src_h, src_w = image_src.shape[:2]
    # print(src_h, src_w)
    dst_h, dst_w = dst_size

    # 判断应该按哪个边做等比缩放
    h = dst_w * (float(src_h) / src_w)  # 按照ｗ做等比缩放
    w = dst_h * (float(src_w) / src_h)  # 按照h做等比缩放

    h = int(h)
    w = int(w)

    if h <= dst_h:
        image_dst = cv2.resize(image_src, (dst_w, int(h)))
    else:
        image_dst = cv2.resize(image_src, (int(w), dst_h))

    h_, w_ = image_dst.shape[:2]
    # print(h_, w_)

    top = int((dst_h - h_) / 2);
    down = int((dst_h - h_ + 1) / 2);
    left = int((dst_w - w_) / 2);
    right = int((dst_w - w_ + 1) / 2);

    value = [0, 0, 0]
    borderType = cv2.BORDER_CONSTANT
    # print(top, down, left, right)
    image_dst = cv2.copyMakeBorder(image_dst, top, down, left, right, borderType, None, value)

    return image_dst

def video_fusion(config,scene):
    video_name = os.path.join(config['root_path'], config['video'], scene + '_' + config['video_name'] + '.mp4')
    if not os.path.exists(os.path.join(config['root_path'], config['video'])):
        os.makedirs(os.path.join(config['root_path'], config['video']))
    # size = tuple(config['video_size'])  # 需要转为视频的图片的尺寸
    fps = config['fps']

    points_image_path = os.path.join(config['root_path'], scene, config['scene'][scene]['out_picture'])
    camera_image_path = os.path.join(config['root_path'], scene, config['scene'][scene]['camera'])

    points_images = os.listdir(points_image_path)
    camera_images = os.listdir(camera_image_path)

    points_images = sorted(points_images)

    points_image = os.path.join(points_image_path, points_images[0])
    camera_image = os.path.join(camera_image_path, points_images[0])

    frame2 = cv2.imread(points_image)
    size = (frame2.shape[1], frame2.shape[0])
    videoWriter = cv2.VideoWriter(video_name, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, size)
    for points_image_ in points_images:
        points_image = os.path.join(points_image_path, points_image_)
        camera_image = os.path.join(camera_image_path, points_image_)
        print(points_image, camera_image)
        frame2 = cv2.imread(points_image)
        frame1 = cv2.imread(camera_image)
        frame1_temp = resize_keep_aspectratio(frame1, tuple(config['image_size']))
        frame2[0:config['image_size'][0], frame2.shape[1] - config['image_size'][1]:frame2.shape[1], :] = frame1_temp
        videoWriter.write(frame2)
    videoWriter.release()
    print('video_name:', video_name)

def create_image(config,scene):
    pr_data_path = os.path.join(config['root_path'], scene, config['scene'][scene]['pr_data_path'])
    pr_data_type = config['scene'][scene]['pr_data_type']
    gt_data_path = os.path.join(config['root_path'], scene, config['scene'][scene]['gt_data_path'])
    gt_data_type = config['scene'][scene]['gt_data_type']
    pts_data_path = os.path.join(config['root_path'], scene, config['scene'][scene]['points'])
    lidar_image_path = os.path.join(config['root_path'], scene, config['scene'][scene]['out_picture'])
    if not os.path.exists(lidar_image_path):
        os.makedirs(lidar_image_path)
    pr_data_list = os.listdir(pr_data_path)
    gt_data_list = os.listdir(gt_data_path)
    pts_data_list = os.listdir(pts_data_path)
    pr_data_list = [i[:-5] for i in pr_data_list]
    gt_data_list = [i[:-5] for i in gt_data_list]
    pts_data_list = [i[:-4] for i in pts_data_list]
    pts_list = np.intersect1d(pr_data_list, gt_data_list, pts_data_list)
    pts_list = sorted(pts_list)

    count = 0
    for file_name in pts_list:
        count += 1

        pr_label_json = os.path.join(pr_data_path, file_name + '.json')
        gt_label_json = os.path.join(gt_data_path, file_name + '.json')
        pts_data_file = os.path.join(pts_data_path, file_name + '.bin')
        lidar_image_file = os.path.join(lidar_image_path, file_name + '.jpg')

        if pr_data_type == 'json':
            name_pr, id_pr, box_pr, velocity_pr, acceleration_pr, score = get_json_data(pr_label_json)
        elif pr_data_type == 'jsons':
            name_pr, id_pr, box_pr, velocity_pr, acceleration_pr, score = get_jsons_data(pr_label_json)
        else:
            name_pr, id_pr, box_pr, velocity_pr, acceleration_pr, score = []
        if gt_data_type == 'json':
            name_gt, id_gt, box_gt, velocity_gt, acceleration_gt = get_json_data(gt_label_json)
        elif gt_data_type == 'jsons':
            name_gt, id_gt, box_gt, velocity_gt, acceleration_gt = get_jsons_data(gt_label_json)
        else:
            name_gt, id_gt, box_gt, velocity_gt, acceleration_gt = []

        name_pr = ['pr_' + str(i) for i in id_pr]
        name_gt = ['gt_' + str(i) for i in id_gt]
        # box_gt = np.array(box_gt)
        # box_gt[:,2]=box_gt[:,2]+3.15
        # box_gt = box_gt.tolist()

        name = name_gt + name_pr
        box = np.array(box_gt + box_pr)

        # pass
        points = get_lidar(pts_data_file)

        if config['bev_flag']:
            points[:, 2] = 0
            box[:, 2] = 0
            box[:, 5] = 1e-10

        start_time = time.time()

        fig = mlab.figure(figure=None, bgcolor=config['color_maps']['bgcolor'], fgcolor=config['color_maps']['fgcolor'],
                          engine=None, size=config['size'])
        # 绘制点云
        fig = visualize_pts(points, show_intensity=False, color=config['color_maps']['points'], size=config['size'],
                            fig=fig)
        # 绘制目标
        fig = draw_label3d(box=box, name=name, config=config, fig=fig)
        mlab.view(azimuth=config['scene_param']['azimuth'], elevation=config['scene_param']['elevation'],
                  distance=config['scene_param']['distance'], roll=config['scene_param']['roll'],
                  focalpoint=config['scene_param']['focalpoint'], figure=fig)
        # 看效果
        mlab.show(stop=True)
        # mlab.savefig(lidar_image_file, figure=fig)
        # mlab.close(all=True)
        end_time = time.time()
        print(':elapsed time is', end_time - start_time, '(s)')

if __name__ == '__main__':
    config_file = 'config/lidar_video_config.yaml'
    with open(config_file, 'r') as f:
        config = yaml.full_load(f)
    for keys in config['color_maps']:
        if keys=='classes':
            for type in config['color_maps']['classes']:
                config['color_maps']['classes'][type]=tuple(config['color_maps']['classes'][type])
        else:
            config['color_maps'][keys]=tuple(config['color_maps'][keys])
    for scene in config['scene']:
        # 渲染图片
        create_image(config, scene)
        # 生成视频
        video_fusion(config,scene)
